import numpy as np
import mindspore as ms
import mindspore.nn as nn
from mindspore.ops import functional as F
from mindspore.common.initializer import initializer


class TestNet(nn.Cell):
    def __init__(self):
        super(TestNet, self).__init__()
        self.op = nn.LSTM(10, 20, 2)

    def construct(self, x, h0, c0):
        out, (hn, cn) = self.op(x, (h0, c0))
        return [out]


if __name__ == "__main__":
    net = TestNet()
    x = ms.Tensor(np.ones([5, 3, 10]).astype(np.float32), ms.float32)
    h0 = ms.Tensor(np.ones([2, 3, 20]).astype(np.float32), ms.float32)
    c0 = ms.Tensor(np.ones([2, 3, 20]).astype(np.float32), ms.float32)
    result = net(x, h0, c0)
    print(result[0])
    print(result[0].mean())
    print(result[0].shape)
